#!/usr/bin/env python
# -*- coding: utf-8 -*-

""" By Martin Senande-Rivera
	For Spatial and temporal expansion of global wildland fire activity in response to climate change """

import numpy as np
import xarray as xr

### File reading ###
####################

ruta_gfed4='../../DATA/GFED4/'	# GFED4 data path
ds_ba_a = xr.open_dataset(ruta_gfed4+'GFED4_BA_a_mean.nc')
ds_ba_m = xr.open_dataset(ruta_gfed4+'GFED4_BA_m_mean.nc')
ds_ba_years = xr.open_dataset(ruta_gfed4+'GFED4_BA_years.nc')
ds_ba_months = xr.open_dataset(ruta_gfed4+'GFED4_BA_months.nc')

BA_years = ds_ba_years.BurnedArea	# Burned Area by year
BA_months = ds_ba_months.BurnedArea	# Burned Area by month

# We consider locations as fire-prone if annual BA>100ha
BA_thres = 100.
BA_y = xr.where(BA_years<BA_thres,np.nan,BA_years)
BA_m = BA_months.copy()
for i in range(BA_years.year.size):		# Loop over the years
	year = BA_years.year[i].values		# Year value
	BA_m.values[12*i:12+12*i,:,:] = xr.where(BA_years.sel(year=year).values<BA_thres,np.nan,BA_m.values[12*i:12+12*i,:,:])


BA_a = BA_y.mean(dim='year',skipna=True)
BA_m = BA_m.groupby('time.month').mean(dim='time',skipna=True)
BA_a.values[np.isnan(BA_a.values)]=0.
BA_m.values[np.isnan(BA_m.values)]=0.

### Monthly Summations ###
##########################
BA_months_max = xr.full_like(BA_m, fill_value=0.)		# For each spatial point, this array will contain 12 values
														# The 1st value is the burned area of the month with maximum burned area
														# The 2nd value is the total burned area of the 2 months with maximum burned area
														# The ith value is the total burned area of the i months with maximum burned area
BA_months_argmax = xr.full_like(BA_m, fill_value=0.)	# For each spatial point, this array will contain 12 values
														# The 1st value is the index of the month with maximum burned area
														# The 2nd value is the index of 2nd month with maximum burned area
														# The ith value is the index of i-th month with maximum burned area
BA_tot = BA_m.copy()
SUM_max=0.
for i in range(0,12):									# Loop along months
	SUM_max = SUM_max+BA_tot.max(dim='month',skipna=True).copy()	# Summation of the burned area of the i months with maximum burned area
	BA_months_max.sel(month=i+1).values[:] = SUM_max.values[:].copy()
	BA_months_argmax.sel(month=i+1).values[:] = BA_tot.argmax(dim='month',skipna=True,keep_attrs=True).copy()+1
	BA_tot[BA_tot.argmax(dim='month',skipna=True,keep_attrs=True)]=0. 	# We set the i-the maximum to 0 to find the i+1 maximum in the next step

### Fire Season Length (FSL) calculation ###
############################################
BA_80 = BA_a.values*0.8
FSL = xr.full_like(BA_a, fill_value=0.)
FSL = xr.where(BA_months_max.sel(month=12)>BA_80,12.,FSL)	# Where the summation over 12 months is greater than the mean annual BA, the FSL is 12
FSL = xr.where(BA_months_max.sel(month=11)>BA_80,11.,FSL)	# Where the summation over 11 months is greater than the mean annual BA, the FSL is 11
FSL = xr.where(BA_months_max.sel(month=10)>BA_80,10.,FSL)	# Where the summation over 10 months is greater than the mean annual BA, the FSL is 10
FSL = xr.where(BA_months_max.sel(month=9)>BA_80,9.,FSL)	# Where the summation over 9 months is greater than the mean annual BA, the FSL is 9
FSL = xr.where(BA_months_max.sel(month=8)>BA_80,8.,FSL)	# Where the summation over 8 months is greater than the mean annual BA, the FSL is 8
FSL = xr.where(BA_months_max.sel(month=7)>BA_80,7.,FSL)	# Where the summation over 7 months is greater than the mean annual BA, the FSL is 7
FSL = xr.where(BA_months_max.sel(month=6)>BA_80,6.,FSL)	# ...
FSL = xr.where(BA_months_max.sel(month=5)>BA_80,5.,FSL)	# ...
FSL = xr.where(BA_months_max.sel(month=4)>BA_80,4.,FSL)	# ...
FSL = xr.where(BA_months_max.sel(month=3)>BA_80,3.,FSL)	# Each line will replace the values of the previous line
FSL = xr.where(BA_months_max.sel(month=2)>BA_80,2.,FSL)	# So we must start in 12 and end in 1
FSL = xr.where(BA_months_max.sel(month=1)>BA_80,1.,FSL)	# The FSL is therefore the minimum number of months with a BA greater than the annual BA

### First Season array (FS_array) calculation ###
#################################################
# Here we compute a 3D array with dimensions (monts,latitude,longitude) filled with zeros and ones.
# The months that are included in the Fire Season of each spatial point are filled with 1.
# The months that are not included in the Fire Season of each spatial point are filled with 0.
FS_array = xr.full_like(BA_m, fill_value=0.)
FSM = BA_months_argmax.copy()	# Array with the list of months sorted by BA for each spatial point.
for i in range(1,13):			# FS_array months loop
	FS_array.values[i-1,:,:][(FSL.values>=1.) & (FSM.sel(month=1).values==i)]=1. 	# i month is in the FS if it is the max BA month (FSM.sel(month=1)) and the FSL is >=1.
	FS_array.values[i-1,:,:][(FSL.values>=2.) & (FSM.sel(month=2).values==i)]=1. 	# i month is in the FS if it is the 2nd max BA month (FSM.sel(month=2)) and the FSL is >=2.
	FS_array.values[i-1,:,:][(FSL.values>=3.) & (FSM.sel(month=3).values==i)]=1. 	# i month is in the FS if it is the 3rd max BA month (FSM.sel(month=3)) and the FSL is >=2.
	FS_array.values[i-1,:,:][(FSL.values>=4.) & (FSM.sel(month=4).values==i)]=1. 	# ...
	FS_array.values[i-1,:,:][(FSL.values>=5.) & (FSM.sel(month=5).values==i)]=1. 	# ...
	FS_array.values[i-1,:,:][(FSL.values>=6.) & (FSM.sel(month=6).values==i)]=1. 	# ...
	FS_array.values[i-1,:,:][(FSL.values>=7.) & (FSM.sel(month=7).values==i)]=1. 	# ...
	FS_array.values[i-1,:,:][(FSL.values>=8.) & (FSM.sel(month=8).values==i)]=1. 	# ...
	FS_array.values[i-1,:,:][(FSL.values>=9.) & (FSM.sel(month=9).values==i)]=1. 	# ...
	FS_array.values[i-1,:,:][(FSL.values>=10.) & (FSM.sel(month=10).values==i)]=1. 	# ...
	FS_array.values[i-1,:,:][(FSL.values>=11.) & (FSM.sel(month=11).values==i)]=1. 	# ...
	FS_array.values[i-1,:,:][(FSL.values>=12.) & (FSM.sel(month=12).values==i)]=1. 	# i month is in the FS if it is the min BA month (FSM.sel(month=12)) and the FSL is =12.

### Save files ###
##################

# FS_array - Binary array over 12 months with 1 in fire season months and 0 in non fire season months. 									Shape: time-12,lat-720,lon-1440
# FS_fullarray - Binary array over 12 months with 1 in fire season months and 0 in non fire season months copied for all years.			Shape: time-252,lat-720,lon-1440

FS_array.name='FS_array'
FS_array.to_netcdf('FS_array.nc')

FS_fullarray = xr.full_like(BA_months, fill_value=0.)
for i in range(BA_years.year.size):
	FS_fullarray.values[12*i:12+12*i,:,:]=FS_array.values[:]

FS_fullarray.name='FS_array'
FS_fullarray.to_netcdf('FS_fullarray.nc')




